'''
Created on Nov 8, 2011

@author: oabalbin
'''
import sys
import numpy as np
import networkx as nx
from collections import defaultdict
import interactome.build_interactome.extract_network as it

def read_ids(nfile):
    ifile=open(nfile)
    h=ifile.next().strip('\r\n').split('\t')
    gids=defaultdict(list)
    i=0
    for l in ifile:
        f=l.strip('\r\n').split('\t')
        gids[f[0]].append(i)
        i+=1
    ifile.close()
    return gids,h

def read_edges(edgefile):
    ifile=open(edgefile)
    edges=[]
    G=nx.Graph()
    for l in ifile:
        if l.startswith('#'):
            continue
        f=l.strip('\n').split('\t')
        G.add_edge(f[0], f[1], weight=f[2])
        edges.append(f[:1])
    ifile.close()
    return G

def read_plain_edges(edgefile):
    ifile=open(edgefile)
    edges=[]
    for l in ifile:
        if l.startswith('#'):
            continue
        f=l.strip('\n').split('\t')
        edges.append([f[0], f[1]])
    ifile.close()
    return edges

def get_edges_weight(edgeslist):
    '''
    get the weight of edges from string database
    myInt is an myInteractome class from the Interactome package
    
    interaction description fields
    self.item_id_a, self.item_id_b, self.neighborhood, self.fusion, self.cooccurence, self.coexpression, self.experimental, \
        self.indatabase, self.textmining, self.combined_score, self.item_id_a_hugo, self.item_id_b_hugo
    
    Define the score to use. By default use string8 combined score
    '''
    myInt = it.myInteractome()
    # Create an empty network
    STRING_max_score = 1000.0
    no_edge_weight= 1.0
    G=nx.Graph()
    
    for edge in edgeslist:
        ta,ib = edge[0],edge[1]
        # returns list of proteins in ensembl nomenclature
        prots, protsrevIndex, notfound,prot2gene_table = myInt.get_proteins_in_network(edge)
        if len(notfound) > 0:
            print "Gene not found in Ensembl=", notfound
            G.add_edge(ta,ib, weight = no_edge_weight)
        else:
            # prots is an dictionary that stores an edge: gene1:prot, gene2:prot
            interact = myInt.STRING.extract_single_p2pinteraction(prots[edge[0]], prots[edge[1]], edge[0], edge[1])
            if interact is not None:
                # Insert the edge in the network
                #edge_weight= 1.0 - float(interact.experimental)/STRING_max_score
                # combined score between experimental and curated database evidence
                #edge_weight= 1.0 - float(interact.experimental)/STRING_max_score * 1.0 - float(interact.indatabase)/STRING_max_score])) # interact.combined_score, combined score
                # edge_weight= 1.0 - float(interact.indatabase)/STRING_max_score
                scores_list = [interact.indatabase]
                edge_weight = calc_STRING_combined_score(scores_list)
                
                G.add_edge(ta,ib, weight = edge_weight)
    return G


def calc_STRING_combined_score(scores_list):
    '''
    It calculatares the combine scores assuming a naive bayes approach in which source of information are independent
    The actual probability should be S =  1 - multplication (1-Si)
    Here we calculate rhe secon term because we want to lowest cost to the high confident interactions.
    '''
    edge_weight=1
    STRING_max_score = 1000.0
    for score in scores_list:
        score_i = 1.0 - float(score)/STRING_max_score
        edge_weight=edge_weight*score_i
    
    return edge_weight
    

def read_nodePrizes(prizesFile):
    '''
    It reads a file nodes transcript_lfc1, protein_lfc2, phospho_lfc3
    '''
    all_nodes=defaultdict()
    ifile=open(prizesFile)
    nodes=defaultdict()
    for l in ifile:
        f=l.strip('\n').split('\t')
        nodes[f[0]]=f[1:]

    ifile.close()
    return nodes


def read_nodePrizes_mat(prizesFile,fcol):
    '''
    It reads a file nodes transcript_lfc1, protein_lfc2, phospho_lfc3
    '''
    gids,h=read_ids(prizesFile)
    mat = np.loadtxt(nfile,skiprows=1,usecols=set(range(fcol,len(h) )))
    # Calculate Maximum across each data set to normalize btw 0-1
    data_max=np.max(mat,axis=0)
    mat = np.true_divide(mat,data_max)
    # Sum the normalize values across data sets
    gene_max=np.sum(mat,axis=1)
    #mat = np.append(mat,np.reshape(fd,(mat.shape[0],-1)),axis=1)
    # create a node dictionary
    nodes=defaultdict()
    for g,i in gids.iteritems():
        nodes[g]=gene_max[i]
    
    return nodes


def map_node2Prizes1(netG,nodePrizes):
    '''
    networkNodes = list of lists
    nodePrizes = dictionary
    It represent each gene as 3 entities: mrna, prot, phospho
    '''
    l=['m','','pp']
    m=['0.0','0.0','0.0']
    nodes=netG.nodes()
    nodes2=defaultdict()
    
    for n in nodes:
        try:
            p=nodePrizes[n]
            for i,j in zip(l,p):
                nodes2[n+i]=str(abs(float( np.power(float(j),2)) )) #NOTE: THIS IS TEMPORAL
                print n+i,j
            # Improve this representation. Here I assume that the interaction is at the level of prot Only
            # This can be inprove by obtaining more info from the kegg map. or the database that gives infor about prot-phosphoprotein interaction
            netG.add_edge(n+'m', n, weight=1)
            netG.add_edge(n, n+'pp', weight=1)
        except KeyError:
            for i,j in zip(l,m):
                nodes2[n+i]=j
                print n,'_'+i,j            
            netG.add_edge(n+'m', n, weight=1)
            netG.add_edge(n, n+'pp', weight=1)
            
    return netG, nodes2


def map_node2Prizes2(netG,nodePrizes):
    '''
    networkNodes = list of lists
    nodePrizes = dictionary
    It represent each gene as 3 entities: mrna, prot, phospho
    '''
    l=['m','','pp']
    mnull='0.0'
    nodes=netG.nodes()
    nodes2=defaultdict()
    
    for n in nodes:
        try:
            p=str(nodePrizes[n][0])
            nodes2[n]=p
        except KeyError:
            nodes2[n]=mnull
            
    return netG, nodes2


def main(edgesFile, nodesPrizeFile):
    '''
    '''
    edges=read_edges(edgesFile)
    nodes=read_nodePrizes(nodesPrizeFile)
    ofed=open(edgesFile+'_heinz','w')
    ofmd=open(nodesPrizeFile+'_heinz','w')
    # return an extend edges list and a nodePrize list
    netG,nodesPrizes = map_node2Prizes1(edges,nodes)
    
    for e in netG.edges(data=True):
        ofed.write(",".join(map(str,[e[0],e[1],e[2]['weight']])).replace(',','\t')+'\n')
        
    for n,v in nodesPrizes.iteritems():
        ofmd.write(",".join([n,v]).replace(',','\t')+'\n')
    
    ofed.close()
    ofmd.close()

def main2(edgesFile, nodesPrizeFile):
    '''
    '''
    fcol=1
    # returns a network
    tmepG=read_edges(edgesFile)
    nodes=read_nodePrizes_mat(nodesPrizeFile,fcol)
    ofed=open(edgesFile+'_heinz2','w')
    ofmd=open(nodesPrizeFile+'_heinz2','w')
    # return an extend edges list and a nodePrize list
    netG,nodesPrizes = map_node2Prizes2(tmepG,nodes)
    ofed.write('#nodeA\tnodeB\tweight\n')
    for e in netG.edges(data=True):
        ofed.write(",".join(map(str,[e[0],e[1],e[2]['weight']])).replace(',','\t')+'\n')
    ofmd.write('#nodeA\tweight\n')
    for n,v in nodesPrizes.iteritems():
        ofmd.write(",".join([n,v]).replace(',','\t')+'\n')
    
    ofed.close()
    ofmd.close()


def main3(edgesFile, nodesPrizeFile):
    '''
    '''
    fcol=1
    edges=read_plain_edges(edgesFile)
    # Return a graph
    tmpG = get_edges_weight(edges)
    nodes=read_nodePrizes_mat(nodesPrizeFile,fcol)
    #sys.exit(0)
    ofed=open(edgesFile+'_heinz2','w')
    ofmd=open(nodesPrizeFile+'_heinz2','w')
    # return an extend edges list and a nodePrize list
    netG,nodesPrizes = map_node2Prizes2(tmpG,nodes)
    ofed.write('#nodeA\tnodeB\tweight\n')
    for e in netG.edges(data=True):
        ofed.write(",".join(map(str,[e[0],e[1],e[2]['weight']])).replace(',','\t')+'\n')
    ofmd.write('#nodeA\tweight\n')
    for n,v in nodesPrizes.iteritems():
        ofmd.write(",".join([n,v]).replace(',','\t')+'\n')
    
    ofed.close()
    ofmd.close()



efile='/data/projects/pcst/test_2011/hsaMerged_edges_hr.tab'
nfile='/data/projects/pcst/test_2011/foldChange.tab'
#main(efile, nfile)
#main2(efile, nfile)
main3(efile, nfile)